# 专门负责餐厅推荐的Agent
from python_a2a import A2AServer, run_server, TaskStatus, TaskState, AgentCard, AgentSkill
import json

class RestaurantAgent(A2AServer):
    def __init__(self, agent_card):
        super().__init__(agent_card=agent_card)
        # 预定义的餐厅数据库
        self.restaurants_db = {
            "beijing": {
                "景点周边": {
                    "天安门广场": [
                        {"name": "全聚德(前门店)", "type": "北京烤鸭", "price": "高档", "distance": "500米"},
                        {"name": "东来顺(王府井店)", "type": "涮羊肉", "price": "中档", "distance": "1公里"}
                    ],
                    "故宫博物院": [
                        {"name": "四季民福", "type": "北京菜", "price": "中高档", "distance": "800米"},
                        {"name": "庆丰包子铺", "type": "包子小吃", "price": "平价", "distance": "300米"}
                    ],
                    "颐和园": [
                        {"name": "听鹂馆", "type": "宫廷菜", "price": "高档", "distance": "园内"},
                        {"name": "万寿山餐厅", "type": "家常菜", "price": "中档", "distance": "园内"}
                    ]
                },
                "天气相关": {
                    "寒冷天气": [
                        {"name": "海底捞", "type": "火锅", "price": "中档", "reason": "暖身御寒"},
                        {"name": "小肥羊", "type": "涮羊肉", "price": "中档", "reason": "温补暖胃"}
                    ],
                    "炎热天气": [
                        {"name": "西贝莜面村", "type": "西北菜", "price": "中档", "reason": "清爽不油腻"},
                        {"name": "新辣道", "type": "川菜", "price": "中档", "reason": "开胃下饭"}
                    ]
                }
            },
            "shanghai": {
                "景点周边": {
                    "外滩": [
                        {"name": "老正兴", "type": "本帮菜", "price": "高档", "distance": "500米"},
                        {"name": "南翔馒头店", "type": "小笼包", "price": "平价", "distance": "2公里"}
                    ]
                }
            }
        }
        print(f"🍽️ RestaurantAgent 初始化完成")

    def handle_task(self, task):
        """处理餐厅推荐任务"""
        message_data = task.message or {}
        content = message_data.get("content", {})
        text = content.get("text", "")
        
        print(f"🍽️ RestaurantAgent收到任务: {text}")
        
        # 解析来自TravelAgent的信息
        request_data = json.loads(text) if text.startswith('{') else {"query": text}
        
        recommendations = []
        
        if "recommendations" in request_data:
            # 来自TravelAgent的景点推荐
            city = request_data.get("city", "beijing")
            attractions = request_data["recommendations"]
            weather_info = request_data.get("weather_info", {})
            
            # 根据景点推荐餐厅
            for attraction in attractions:
                attraction_name = attraction["name"]
                nearby_restaurants = self._get_restaurants_near_attraction(city, attraction_name)
                recommendations.extend(nearby_restaurants)
            
            # 根据天气推荐特殊餐厅
            weather_restaurants = self._get_restaurants_by_weather(city, weather_info)
            recommendations.extend(weather_restaurants)
            
            response = {
                "agent": "RestaurantAgent",
                "city": city,
                "based_on_attractions": [attr["name"] for attr in attractions],
                "restaurant_recommendations": recommendations,
                "total_restaurants": len(recommendations)
            }
        else:
            # 直接餐厅查询
            city = request_data.get("city", "beijing")
            recommendations = self._get_popular_restaurants(city)
            
            response = {
                "agent": "RestaurantAgent",
                "city": city,
                "based_on_attractions": [],
                "restaurant_recommendations": recommendations,
                "note": "通用餐厅推荐，建议结合景点行程获得更精准推荐"
            }

        task.artifacts = [{"parts": [{"type": "text", "text": json.dumps(response, ensure_ascii=False)}]}]
        task.status = TaskStatus(state=TaskState.COMPLETED)
        print(f"🍽️ RestaurantAgent回复: 推荐了{len(recommendations)}家餐厅")
        
        # 发送数据流信息到UI
        try:
            import requests
            requests.post("http://localhost:8000/api/notify", json={
                "type": "data_flow",
                "step": {
                    "agent": "RestaurantAgent",
                    "description": f"餐厅推荐完成 - {city}",
                    "input": request_data,
                    "output": response,
                    "upstream": "CoordinatorAgent",
                    "downstream": "CoordinatorAgent"
                }
            }, timeout=2)
        except:
            pass
        
        return task

    def _get_restaurants_near_attraction(self, city, attraction_name):
        """获取景点附近的餐厅"""
        city_data = self.restaurants_db.get(city, {})
        attraction_restaurants = city_data.get("景点周边", {})
        
        # 精确匹配景点名称
        if attraction_name in attraction_restaurants:
            restaurants = attraction_restaurants[attraction_name]
            for restaurant in restaurants:
                restaurant["location_reason"] = f"靠近{attraction_name}"
            return restaurants
        
        # 模糊匹配
        for key, restaurants in attraction_restaurants.items():
            if any(word in attraction_name for word in key.split()) or any(word in key for word in attraction_name.split()):
                for restaurant in restaurants:
                    restaurant["location_reason"] = f"靠近{key}区域"
                return restaurants[:2]  # 返回前2家
                
        return []

    def _get_restaurants_by_weather(self, city, weather_info):
        """根据天气推荐餐厅"""
        restaurants = []
        city_data = self.restaurants_db.get(city, {})
        weather_restaurants = city_data.get("天气相关", {})
        
        temperature = weather_info.get("temperature_info", {}).get("temp")
        
        if temperature and temperature < 10:
            # 寒冷天气
            cold_restaurants = weather_restaurants.get("寒冷天气", [])
            restaurants.extend(cold_restaurants[:1])  # 取1家
        elif temperature and temperature > 25:
            # 炎热天气
            hot_restaurants = weather_restaurants.get("炎热天气", [])
            restaurants.extend(hot_restaurants[:1])  # 取1家
            
        return restaurants

    def _get_popular_restaurants(self, city):
        """获取热门餐厅"""
        city_data = self.restaurants_db.get(city, {})
        popular = []
        
        # 从各个景点取一些代表性餐厅
        attraction_restaurants = city_data.get("景点周边", {})
        for attraction, restaurants in list(attraction_restaurants.items())[:2]:
            if restaurants:
                restaurant = restaurants[0].copy()
                restaurant["popularity_reason"] = "热门景点附近"
                popular.append(restaurant)
                
        return popular

if __name__ == "__main__":
    agent_card = AgentCard(
        name="Restaurant Expert Agent", 
        description="专业餐厅推荐Agent，基于景点位置和天气条件推荐合适的餐厅",
        url="http://localhost:7005",
        version="1.0.0",
        skills=[
            AgentSkill(
                name="Location-based Restaurant Recommendation",
                description="根据景点位置推荐附近餐厅"
            ),
            AgentSkill(
                name="Weather-based Cuisine Recommendation", 
                description="根据天气条件推荐合适的菜系"
            ),
            AgentSkill(
                name="Dining Plan Integration",
                description="整合用餐计划到旅游行程中"
            )
        ]
    )

    restaurant_agent = RestaurantAgent(agent_card)
    
    print("🍽️ RestaurantAgent 启动于 http://localhost:7005")
    run_server(restaurant_agent, host="0.0.0.0", port=7005)
